import encoding
import threading
import SocketServer
import socket
import struct
import os
import time
from Queue import Queue
from Crypto.Random import random
from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
from Crypto.Hash import RIPEMD
from Crypto.Hash import HMAC
from enigma import EnigmaMachine, add_uppercase

listen_at = ('127.0.0.1', 4500)
server_addr = ('127.0.0.1', 3967)
RSA_keyfile = '.\pubkey.key'
enigma_settings = {'count': 2, 'rotors': 'II VI V', 'ring_settings': "1 16 7", 'reflector': 'B'}
complete_alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
add_padding = True

class Logger(object):
    def __init__(self):
        self.lock = threading.Lock()

    def log(self, str):
        self.lock.acquire()
        print(str)
        self.lock.release()

logger = Logger()

logger.log('Starting...')

EnigmaMachine.wrapper_initialize(**enigma_settings)

try:
    server_socket = socket.create_connection(server_addr)
    seedstr = server_socket.recv(72)[0:4]
    if not seedstr:
        logger.log('Connection closed.')
        os._exit(0)
except:
    logger.log('Failed to connect to server.')
    os._exit(0)

enigmakey =''.join(random.sample(complete_alphabet, 3))
hmackey = struct.pack("!I", random.getrandbits(32))

alphabet = ''.join(random.sample(complete_alphabet, 24))

with open(RSA_keyfile, 'r') as RSA_key:
    key = RSA.importKey(RSA_key.read())
    rsa = PKCS1_OAEP.new(key)
padding_count = random.randint(0, 16)
padding = ''
for i in xrange(padding_count):
    padding += struct.pack("!I", random.getrandbits(32))
server_socket.send(rsa.encrypt(seedstr + enigmakey + alphabet + hmackey) + padding)

upstream_hmac = HMAC.new(hmackey, digestmod=RIPEMD.new()) # used to sign data waiting to be sent
downstream_hmac = HMAC.new(hmackey, digestmod=RIPEMD.new()) # used to verify data we have received
upstream_enigma = EnigmaMachine() # used to encrypt data waiting to be sent
upstream_enigma.set_display(enigmakey)
downstream_enigma = EnigmaMachine() # used to decrypt data we have received
downstream_enigma.set_display(enigmakey)

vf = server_socket.recv(256).replace(' ', '').upper()
try:
    _alphabet = upstream_enigma.process_text(vf)[0:24]
except:
    logger.log('Received bad handshake data from server.')
    os._exit(0)

if not _alphabet == alphabet:
    logger.log('Enigma machine desync.')
    os._exit(0)

logger.log('Connected!')

outbuf = Queue()
clean_conn = Queue()
channels = []

class SenderThread(threading.Thread):
    def run(self):
        while True:
            msg = outbuf.get()
            downstream_hmac.update(msg)
            hmac = downstream_hmac.digest()

            length = struct.pack("!I", len(hmac) + len(msg))

            encdata = encoding.add_space(add_uppercase(downstream_enigma.process_text(encoding.enbase(random.choice(alphabet) + length + hmac + msg, alphabet))))

            server_socket.sendall(encdata)

class ReceiverThread(threading.Thread):
    def run(self):
        inbuf = []
        ib = ''
        while True:
            data = server_socket.recv(65536)
            if not data:
                logger.log('Connection closed.')
                os._exit(0)
            try:
                data = ib + upstream_enigma.process_text(data.replace(' ', ''))
                ib = ''
                if not len(data) % 2 == 0:
                    data, ib = data[0:-1], data[-1]
                inbuf.append(encoding.debase(data, alphabet))
            except:
                logger.log('Exception decoding data from server.')
                os._exit(0)

            if inbuf == []:
                continue

            idata = ''.join(inbuf)
            inbuf = []

            while True:
                if not idata:
                    break

                try:
                    length, = struct.unpack("!I", idata[1:5])
                except:
                    logger.log('Exception decoding data from server.')
                    os._exit(0)
                if len(idata)-5 < length:
                    inbuf.append(idata)
                    break

                msg, idata = idata[5:5+length], idata[5+length:]

                hmac = msg[0:20]
                upstream_hmac.update(msg[20:])

                if not upstream_hmac.digest() == hmac:
                    logger.log('Got a message with bad HMAC from ', addr)
                    os._exit(0)

                if msg[20] == 'c':
                    clean_conn.put(msg[21:25])
                elif msg[20] == 'd':
                    for i in channels:
                        if i.connected(msg[21:25]):
                            i.disconnect()
                            break
                elif msg[20] == 'r':
                    for i in channels:
                        if i.connected(msg[21:25]):
                            i.push(msg[25:])
                            break
                elif msg[20] == 'p':
                    pass
                else:
                    logger.log('Connection closed by server.')
                    os._exit(0)

class PaddingThread(threading.Thread):
    def run(self):
        while True:
            padding_count = random.randint(0, 42)
            padding = ''
            for i in xrange(padding_count):
                padding += struct.pack("!I", random.getrandbits(32))
            outbuf.put('p' + padding)
            time.sleep(random.randint(0, 42)/10)

class RequestHandler(SocketServer.StreamRequestHandler):
    def handle(self):
        self.disconnected = False
        self.in_queue = Queue()
        self.sender_thread = threading.Thread(target=self.send_to_req)
        self.sender_thread.start()
        outbuf.put('c')
        channels.append(self)
        self.id = clean_conn.get()

        while True:
            if self.disconnected:
                return
            data = self.request.recv(65536)
            if not data:
                outbuf.put('d' + self.id)
                return
            outbuf.put('m' + self.id + data)

    def connected(self, num):
        if self.disconnected or not self.id:
            return False
        else:
            return self.id == num

    def disconnect(self):
        self.disconnected = True

    def push(self, msg):
        self.in_queue.put(msg)

    def send_to_req(self):
        while True:
            msg = self.in_queue.get()
            try:
                self.request.sendall(msg)
            except:
                outbuf.put('d' + self.id)
                self.disconnected = True
                break

if __name__ == '__main__':
    if add_padding:
        PaddingThread().start()
    SenderThread().start()
    ReceiverThread().start()
    SocketServer.ThreadingTCPServer(listen_at, RequestHandler).serve_forever()